import pandas as pd
import numpy as np
from sklearn import metrics
import lightgbm as lgb
from sklearn.model_selection import StratifiedKFold

params = {
    'n_estimators': 5000,
    'early_stopping_rounds': 100,
    'boosting_type': 'gbdt',
    'objective': 'multiclass',
    'num_class': 3
}

ori_data1 = pd.read_csv('./data/X_xyvd_d2_100.csv', header=0)
ori_data2 = pd.read_csv('./data/X_desc.csv', header=0)
ori_data3 = pd.read_csv('./data/trn_per_300.csv', header=0)

tst_data1 = pd.read_csv('./data/X_xyvd_tst_d2_100.csv', header=0)
tst_data2 = pd.read_csv('./data/tst_desc.csv', header=0)
tst_data3 = pd.read_csv('./data/tst_per_300.csv', header=0)
tst_data_ = pd.merge(tst_data1, tst_data2, on='ship', how='left')
tst_data_ = pd.merge(tst_data_, tst_data3, on='ship', how='left')

trn_data1 = ori_data1
trn_data2 = ori_data2.drop('type', axis=1)
trn_data3 = ori_data3.drop('type', axis=1)
trn_data_ = pd.merge(trn_data1, trn_data2, on='ship', how='left')
trn_data_ = pd.merge(trn_data_, trn_data3, on='ship', how='left')


feasA = []
for trn_data in [ori_data1, ori_data2, ori_data3]:
    X = trn_data.iloc[:, :-2]
    y = trn_data['type']

    fold = StratifiedKFold(n_splits=7, shuffle=True, random_state=100)
    for index, (train_idx, val_idx) in enumerate(fold.split(X, y)):
        train_set = lgb.Dataset(X.iloc[train_idx], y.iloc[train_idx])
        val_set = lgb.Dataset(X.iloc[val_idx], y.iloc[val_idx])
        model = lgb.train(params, train_set, valid_sets=[train_set, val_set],
                          verbose_eval=False)

        trn_pred = np.argmax(model.predict(X.iloc[train_idx]), axis=1)
        f1_trn = metrics.f1_score(y.iloc[train_idx], trn_pred, average='macro')

        val_pred = np.argmax(model.predict(X.iloc[val_idx]), axis=1)
        f1_val = metrics.f1_score(y.iloc[val_idx], val_pred, average='macro')
        print(index, 'trn f1:', f1_trn, 'val f1:', f1_val)

        break

    header = pd.DataFrame(X.columns)
    feat_imp = model.feature_importance()
    sort_ind = feat_imp.argsort()[::-1]
    A = pd.concat([header.iloc[sort_ind, :].reset_index().drop('index', axis=1),
                   pd.DataFrame(feat_imp[sort_ind])], axis=1, ignore_index=True)

    feasA += [A]

numfeas = 300
feas_all = []
weights0 = np.array([0.8126011533941019, 0.8886301416470362, 0.9148359432509245])-0.8
for i in range(3):
    A = feasA[i]
    weight_num = np.ceil(weights0/np.sum(weights0)*numfeas).astype(int)
    feas = A.iloc[:weight_num[i], :][0].tolist()
    feas_all += feas

feas_sp = ['x_min', 'x_max', 'y_min', 'y_max', 'v_min', 'v_max', 'd_min', 'd_max', 'hour', 'weekday']
feas_out = [a for a in feas_all if a not in feas_sp] + ['ship', 'type']
trn_out = trn_data_[feas_out]
trn_out.to_csv(f'data/feas_lgb_merge_{numfeas}.csv', header=True, index=False)
feas_tst_out = [a for a in feas_all if a not in feas_sp] + ['ship']
tst_out = tst_data_[feas_tst_out]
tst_out.to_csv(f'data/feas_lgb_merge_tst_{numfeas}.csv', header=True, index=False)